package com.shujia.flink.state

import org.apache.flink.api.common.functions.{ReduceFunction, RuntimeContext}
import org.apache.flink.api.common.state.{ReducingState, ReducingStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

object Demo4ReducingState {
  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    val studentsDS: DataStream[String] = env.socketTextStream("master", 8888)

    /**
      * 实时统计每个班级的平均年龄
      *
      */
    val kvDS: DataStream[(String, Double)] = studentsDS.map(line => {
      val split: Array[String] = line.split(",")

      (split(4), split(2).toDouble)
    })

    kvDS
      .keyBy(_._1)
      .process(new KeyedProcessFunction[String, (String, Double), (String, Double)] {


        var sumAgeState: ReducingState[Double] = _

        var sumNumState: ReducingState[Double] = _

        override def open(parameters: Configuration): Unit = {
          val context: RuntimeContext = getRuntimeContext

          /**
            * ReducingState： 聚合状态，需要一个聚合函数，和valueState类似，也是保存一个单值
            *
            */

          //保存总年龄的状态
          val sumAgeReduceStateDesc = new ReducingStateDescriptor[Double]("sum_age", new ReduceFunction[Double] {
            override def reduce(value1: Double, value2: Double): Double = {
              value1 + value2
            }
          }, classOf[Double])


          sumAgeState = context.getReducingState(sumAgeReduceStateDesc)


          //保存总人数的状态
          val sumNumReduceStateDesc = new ReducingStateDescriptor[Double]("sum_num", new ReduceFunction[Double] {
            override def reduce(value1: Double, value2: Double): Double = {
              value1 + value2
            }
          }, classOf[Double])


          sumNumState = context.getReducingState(sumNumReduceStateDesc)


        }

        override def processElement(value: (String, Double), ctx: KeyedProcessFunction[String, (String, Double), (String, Double)]#Context, out: Collector[(String, Double)]): Unit = {

          //将每一条的年龄保存到存放年龄的状态中
          //调用add 方法，在集合状态的内部会自动累加
          sumAgeState.add(value._2)

          //向总人数的状态中增加一个人
          sumNumState.add(1)


          //计算平均年龄
          val avgAge: Double = sumAgeState.get() / sumNumState.get()

          //将数据发送到下游
          out.collect((value._1, avgAge))


        }
      })
      .print()


    env.execute()
  }
}
